import os
import json
import random
import sys
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

basePath = '../../'

def plot_fov_comparison(movie, fov):
    os.makedirs('./FoV comparison/', exist_ok=True)
    with open(basePath + movie + '/' + movie + '.json', 'r') as load_file:
        movie_json_tmp = json.load(load_file)
        segments = int(len(movie_json_tmp['segment_size_bits']))
    movie_json = np.zeros((6, int(np.max(segments))))
    movie_total = np.zeros(int(np.max(segments)))

    for i in range(segments):
        size_tmp = np.zeros(6)
        size_max_tmp = 0
        # iteration for 6 faces
        for j in range(6):
            size_max_tmp += movie_json_tmp['segment_size_bits'][i][j][3]
            # iteration for 6 possible FoV
            for k in range(6):
                size_tmp[k] += movie_json_tmp['segment_size_bits'][i][j][fov[k][j]]
        for j in range(6):
            movie_json[j][i] = size_tmp[j] / 1000 / 1000
        movie_total[i] = size_max_tmp / 1000 / 1000
    plt.figure(figsize=(15, 10))
    plt.title('Bitrate Record per Segment', fontsize=23)
    plt.plot(range(segments), movie_total, label='Without ABR', color='red')
    plt.plot(range(segments), movie_json[0], label='FoV trace No.0', color='darkorange')
    plt.plot(range(segments), movie_json[1], label='FoV trace No.1', color='gold')
    plt.plot(range(segments), movie_json[2], label='FoV trace No.2', color='mediumseagreen')
    plt.plot(range(segments), movie_json[3], label='FoV trace No.3', color='deepskyblue')
    plt.plot(range(segments), movie_json[4], label='FoV trace No.4', color='darkorchid')
    plt.plot(range(segments), movie_json[5], label='FoV trace No.5', color='grey')
    plt.xlabel('Time (s)', fontsize=17)
    plt.ylabel('Bitrate (Mbps)', fontsize=17)
    if movie == 'BUPT_gate':
        plt.legend(ncol=2, fontsize=12, bbox_to_anchor=(0.4, 0.9), shadow=True, fancybox=True)
    elif movie == 'BUPT_gate_LB':
        plt.legend(ncol=2, fontsize=12, bbox_to_anchor=(0.4, 0.9), shadow=True, fancybox=True)
    elif movie == 'BUPT_road':
        plt.legend(ncol=2, fontsize=12, bbox_to_anchor=(0.95, 0.9), shadow=True, fancybox=True)
    elif movie == 'BUPT_road_LB':
        plt.legend(ncol=2, fontsize=12, bbox_to_anchor=(1, 0.95), shadow=True, fancybox=True)
    elif movie == 'skiing':
        plt.legend(ncol=2, loc='upper right', fontsize=12, bbox_to_anchor=(1, 0.9), shadow=True, fancybox=True)
    sns.despine()
    plt.savefig('./results/FoV comparison/' + movie + '.png')


def plot_percentage(movie, fov):
    path = './FoV comparison in percentage/'
    os.makedirs(path, exist_ok=True)
    with open(basePath + movie + '/' + movie + '.json', 'r') as load_file:
        movie_json_tmp = json.load(load_file)
        segments = int(len(movie_json_tmp['segment_size_bits']))
    # bandwidth = np.zeros((6, segments))
    # expansion = np.zeros((6, segments))
    bandwidth = np.zeros(segments)
    expansion = np.zeros(segments)

    for i in range(segments):
        size_tmp = np.zeros(6)
        size_max_tmp = 0
        # iteration for 6 faces
        for j in range(6):
            size_max_tmp += movie_json_tmp['segment_size_bits'][i][j][3]
            # iteration for 6 possible FoV
            for k in range(6):
                size_tmp[k] += movie_json_tmp['segment_size_bits'][i][j][fov[k][j]]
        bandwidth[i] = (1 - np.max(size_tmp) / size_max_tmp) * 100
        expansion[i] = bandwidth[i] / (1 - bandwidth[i]) * 100
        # bandwidth[i] = np.max(size_tmp) / size_max_tmp * 100
        # for j in range(6):
        #     bandwidth[j][i] = 1 - size_tmp[j] / size_max_tmp
        #     expansion[j][i] = bandwidth[j][i] / (1 - bandwidth[j][i]) * 100
        #     bandwidth[j][i] *= 100

    # 只画出带宽节省
    fig, ax = plt.subplots(figsize=(15, 10))
    ax.plot(range(segments), bandwidth, color='red')
    ax.set_xlabel('Time (s)', fontsize=20)
    ax.set_ylabel('Bandwidth saving (%)', fontsize=20)
    ax.axis([-5, segments + 10, 0, np.max(bandwidth)*1.2])
    ax.spines["top"].set_color("none")
    ax.spines["right"].set_color("none")
    fig.savefig(path + movie + '.png')
    # 一张图画带宽节省/系统扩容 两个曲线重合/对称
    # fig, ax1 = plt.subplots(figsize=(15, 10))
    # ax2 = ax1.twinx()
    # ax1.plot(range(segments), bandwidth, color='red')
    # ax2.plot(range(segments), expansion, color='orange')
    #
    # ax1.set_xlabel('Time (s)', fontsize=17)
    # ax1.set_ylabel('Bandwidth utilization (%)', color='red', fontsize=17)
    # ax2.set_ylabel('System expansion (%)', color='orange', fontsize=17)
    # ax1.spines["top"].set_color("none")
    # ax2.spines["top"].set_color("none")
    # plt.show()

    # 两张图画 6*FoV的带宽节省 / 系统扩容
    # plt.figure(figsize=(15, 10))
    # plt.title('Bandwidth saving', fontsize=23)
    # plt.plot(range(segments), bandwidth[0], label='FoV trace No.0', color='red')
    # plt.plot(range(segments), bandwidth[1], label='FoV trace No.1', color='darkorange')
    # plt.plot(range(segments), bandwidth[2], label='FoV trace No.2', color='gold')
    # plt.plot(range(segments), bandwidth[3], label='FoV trace No.3', color='mediumseagreen')
    # plt.plot(range(segments), bandwidth[4], label='FoV trace No.4', color='deepskyblue')
    # plt.plot(range(segments), bandwidth[5], label='FoV trace No.5', color='darkorchid')
    # plt.xlabel('Time (s)', fontsize=17)
    # plt.ylabel('Percentage (%)', fontsize=17)
    # sns.despine()
    # plt.legend()
    # plt.axis([-5, segments + 10, 0, np.max(bandwidth)*1.2])
    # # print(plt.axis())
    # plt.show()
    # plt.figure(figsize=(15, 10))
    # plt.title('System expansion', fontsize=23)
    # plt.plot(range(segments), expansion[0], label='FoV trace No.0', color='red')
    # plt.plot(range(segments), expansion[1], label='FoV trace No.1', color='darkorange')
    # plt.plot(range(segments), expansion[2], label='FoV trace No.2', color='gold')
    # plt.plot(range(segments), expansion[3], label='FoV trace No.3', color='mediumseagreen')
    # plt.plot(range(segments), expansion[4], label='FoV trace No.4', color='deepskyblue')
    # plt.plot(range(segments), expansion[5], label='FoV trace No.5', color='darkorchid')
    # plt.xlabel('Time (s)', fontsize=17)
    # plt.ylabel('Percentage (%)', fontsize=17)
    # sns.despine()
    # plt.legend()
    # plt.axis([-5, segments + 10, 0, np.max(expansion)*1.2])
    # plt.show()


def plot_movie_comparison(movies, fov):
    os.makedirs('./Movie comparison/', exist_ok=True)
    movie_count = len(movies)
    segments = np.zeros(movie_count)
    for i in range(movie_count):
        with open(basePath + movies[i] + '/' + movies[i] + '.json', 'r') as load_file:
            movie_json_tmp = json.load(load_file)
        segments[i] = len(movie_json_tmp['segment_size_bits'])
    movie_json = np.zeros((movie_count, int(np.max(segments))))
    movie_total = np.zeros((movie_count, int(np.max(segments))))

    for i in range(movie_count):
        with open(basePath + movies[i] + '/' + movies[i] + '.json', 'r') as load_file:
            movie_json_tmp = json.load(load_file)
            for j in range(int(segments[i])):
                size_tmp = 0
                size_max_tmp = 0
                for k in range(6):
                    size_tmp += movie_json_tmp['segment_size_bits'][j][k][fov[k]]
                    size_max_tmp += movie_json_tmp['segment_size_bits'][j][k][3]
                movie_json[i][j] = size_tmp / 1000 / 1000
                movie_total[i][j] = size_max_tmp / 1000 / 1000

    size_withoutABR = np.zeros(movie_count)
    size_withABR = np.zeros(movie_count)

    for i in range(movie_count):
        size_withoutABR[i] = np.sum(movie_total[i]) / segments[i]
        size_withABR[i] = np.sum(movie_json[i]) / segments[i]
    x = list(range(0, 2*movie_count, 2))
    total_width, n = 1.5, 2
    width = total_width / n
    plt.figure('Comparison bitrates between different movies')

    plt.bar(x, size_withABR, width=width, color='orange', label='Buffer-based ABR', align='edge')
    for i in range(movie_count):
        x[i] += width
        if movies[i] == 'BUPT_gate':
            movies[i] = 'Gate HD'
        elif movies[i] == 'BUPT_gate_LB':
            movies[i] = 'Gate SD'
        if movies[i] == 'BUPT_road':
            movies[i] = 'Road HD'
        elif movies[i] == 'BUPT_road_LB':
            movies[i] = 'Road SD'
        elif movies[i] == 'skiing':
            movies[i] = 'Skiing'
    plt.bar(x, size_withoutABR, width=width, color='red', tick_label=movies, label='Highest bitrate', align='edge')

    for x_t, y_t in zip(x, size_withABR):
        plt.text(x_t-0.4, y_t+0.5, '%.2f' % y_t, ha='center', va='bottom')

    for x_t, y_t in zip(x, size_withoutABR):
        plt.text(x_t+0.35, y_t+0.5, '%.2f' % y_t, ha='center', va='bottom')

    plt.title('Comparison in different movies', fontsize=14)
    plt.ylabel('Average bitrate per segment (Mbps)', fontsize=12)
    plt.legend(shadow=True, fancybox=True)
    sns.despine()
    # plt.show()
    plt.savefig('./results/Movie comparison/comparison.png')


def plot_movie_comparison_with_precision(movies, fov, prob):
    path = './Movie comparison with precision/'
    os.makedirs(path, exist_ok=True)
    movie_count = len(movies)
    segments = np.zeros(movie_count)
    for i in range(movie_count):
        with open(basePath + movies[i] + '/' + movies[i] + '.json', 'r') as load_file:
            movie_json_tmp = json.load(load_file)
        segments[i] = len(movie_json_tmp['segment_size_bits'])
    movie_json = np.zeros((movie_count, int(np.max(segments))))
    movie_total = np.zeros((movie_count, int(np.max(segments))))

    for i in range(movie_count):
        with open(basePath + movies[i] + '/' + movies[i] + '.json', 'r') as load_file:
            movie_json_tmp = json.load(load_file)
            for j in range(int(segments[i])):
                size_tmp = 0
                size_max_tmp = 0
                for k in range(6):
                    if random.randint(0, 100) <= prob:
                        size_tmp += movie_json_tmp['segment_size_bits'][j][k][fov[k]]
                    else:
                        size_tmp += movie_json_tmp['segment_size_bits'][j][k][3]
                    size_max_tmp += movie_json_tmp['segment_size_bits'][j][k][3]
                movie_json[i][j] = size_tmp / 1000 / 1000
                movie_total[i][j] = size_max_tmp / 1000 / 1000

    size_withoutABR = np.zeros(movie_count)
    size_withABR = np.zeros(movie_count)

    for i in range(movie_count):
        size_withoutABR[i] = np.sum(movie_total[i]) / segments[i]
        size_withABR[i] = np.sum(movie_json[i]) / segments[i]
    x = list(range(0, 2 * movie_count, 2))
    total_width, n = 1.5, 2
    width = total_width / n
    plt.figure('Comparison bitrates between different movies')

    plt.bar(x, size_withABR, width=width, color='orange', label='Buffer-based ABR', align='edge')
    for i in range(movie_count):
        x[i] += width
        if movies[i] == 'BUPT_gate':
            movies[i] = 'Gate HD'
        elif movies[i] == 'BUPT_gate_LB':
            movies[i] = 'Gate SD'
        if movies[i] == 'BUPT_road':
            movies[i] = 'Road HD'
        elif movies[i] == 'BUPT_road_LB':
            movies[i] = 'Road SD'
        elif movies[i] == 'skiing':
            movies[i] = 'Skiing'
    plt.bar(x, size_withoutABR, width=width, color='red', tick_label=movies, label='Highest bitrate', align='edge')

    for x_t, y_t in zip(x, size_withABR):
        plt.text(x_t - 0.4, y_t + 0.5, '%.2f' % y_t, ha='center', va='bottom')

    for x_t, y_t in zip(x, size_withoutABR):
        plt.text(x_t + 0.35, y_t + 0.5, '%.2f' % y_t, ha='center', va='bottom')

    plt.title('Comparison in different movies', fontsize=14)
    plt.ylabel('Average bitrate per segment (Mbps)', fontsize=12)
    plt.legend(shadow=True, fancybox=True)
    sns.despine()
    # plt.show()
    plt.savefig(path + 'comparison_' + str(prob) + '.png')


if __name__ == '__main__':
    movies = sys.argv[1]
    # plot = sys.argv[2]
    if movies == 'all':
        movies = 'BUPT_gate/BUPT_gate_LB/BUPT_road/BUPT_road_LB/skiing'
    movies = movies.split('/')
    # FoV_1 现在的情况
    fov1 = [
        [3, 0, 1, 2, 2, 1],
        [0, 3, 2, 1, 1, 2],
        [1, 1, 3, 0, 1, 1],
        [1, 1, 0, 3, 1, 1],
        [1, 2, 1, 2, 3, 0],
        [2, 1, 2, 1, 0, 3]
    ]
    # FoV_2 提高1中的选择 2->3 1->2
    fov2 = [
        [3, 0, 2, 3, 3, 2],
        [0, 3, 3, 2, 2, 3],
        [2, 2, 3, 0, 2, 2],
        [2, 2, 0, 3, 2, 2],
        [2, 3, 2, 3, 3, 0],
        [3, 2, 3, 2, 0, 3]
    ]
    # FoV_3 提高1中的选择 1->2
    fov3 = [
        [3, 0, 2, 2, 2, 2],
        [0, 3, 2, 2, 2, 2],
        [2, 2, 3, 0, 2, 2],
        [2, 2, 0, 3, 2, 2],
        [2, 2, 2, 2, 3, 0],
        [2, 2, 2, 2, 0, 3]
    ]
    # FoV_4 提高1中的选择 2->3
    fov4 = [
        [3, 0, 1, 3, 3, 1],
        [0, 3, 3, 1, 1, 3],
        [1, 1, 3, 0, 1, 1],
        [1, 1, 0, 3, 1, 1],
        [1, 3, 1, 3, 3, 0],
        [3, 1, 3, 1, 0, 3]
    ]
    # FoV_5 提高1中的选择 50% 2->3
    fov5 = [
        [3, 0, 1, 3, 2, 1],
        [0, 3, 3, 1, 1, 2],
        [1, 1, 3, 0, 1, 1],
        [1, 1, 0, 3, 1, 1],
        [1, 2, 1, 3, 3, 0],
        [2, 1, 3, 1, 0, 3]
    ]

    # 不同视野消耗带宽对比
    # for i in range(len(movies)):
    #     plot_fov_comparison(movies[i], fov2)

    # 不同视频的平均比特率对比
    # plot_movie_comparison(movies, fov2[0])

    # 添加降低百分比折线图
    # for i in range(len(movies)):
    #     plot_percentage(movies[i], fov5)

    # 考虑预测准确率
    plot_movie_comparison_with_precision(movies, fov5[0], 100)

